#include <iostream>
#include <cassert>
#include <openvino/openvino.hpp>
extern "C" {
#include "stb_image.h"
}

#define NHWC 0

void print_any_value(const ov::Any& value) {
    if (value.empty()) {
        std::cout << "EMPTY VALUE" << std::endl;
    } else {
        std::string stringValue = value.as<std::string>();
        std::cout << (stringValue.empty() ? "\"\"" : stringValue) << std::endl;
    }
}


int main(int argc, char *argv[]) {
    std::cout << ov::get_openvino_version() << std::endl;

    // Initialize OpenVINO Core
    ov::Core core;

    // Get list of available devices
    auto available_devices = core.get_available_devices();
    std::cout << "Available devices:" << std::endl;
    for (const auto& device : available_devices) {
        std::cout << ", " << device << std::endl;
    }

    // Query and print supported metrics and config keys
    for (const auto& device : available_devices) {
        std::cout << "Device: " << device << std::endl;

        // Query supported properties and print all of them
        std::cout << "Supported properties:" << std::endl;
        auto supported_properties = core.get_property(device, ov::supported_properties);
        for (const auto& property : supported_properties) {
            if (property != ov::supported_properties.name()) {
                std::cout << '\t' << (property.is_mutable() ? "Mutable   " : "Read-only ") << property << ": " << std::flush;
                print_any_value(core.get_property(device, property));
            }
        }
    }

    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <model_path> <image_path>" << std::endl;
        return 1;
    }

    auto compiled_model = core.compile_model(argv[1], "CPU");
    std::cout << "Compiled model for device: CPU" << std::endl;
    auto infer_request = compiled_model.create_infer_request();

    auto input_tensor = infer_request.get_input_tensor();
    auto input_shape = input_tensor.get_shape();
    std::cout << "Input tensor shape: " << input_shape << std::endl;
    std::cout << "Input tensor element type: " << input_tensor.get_element_type() << std::endl;

    // Load and set input image
    int width, height, channels;
    unsigned char* image_data = stbi_load(argv[2], &width, &height, &channels, 0);
    std::cout << "Loaded image: " << argv[2] << " with dimensions: " << width << "x" << height << " and channels: " << channels << std::endl;
    assert((width == 28) && (height == 28) && (channels == 1));
    if (!image_data) {
        std::cerr << "Failed to load image: " << argv[2] << std::endl;
        return 1;
    }
    // Convert uint8 to float32 and HWC to CHW format
    float* input_data = new float[height * width * channels];
#if NHWC
    for (int i = 0; i < height * width * channels; ++i) {
        input_data[i] = (static_cast<float>(image_data[i]) / 255.0f - 0.1307f) / 0.3081f; // Normalize to [0, 1] range
    }
#else
    for (int h = 0; h < height; ++h) {
        for (int w = 0; w < width; ++w) {
            for (int c = 0; c < channels; ++c) {
                int index = (h * width + w) * channels + c;
                input_data[c * height * width + h * width + w] = (static_cast<float>(image_data[index]) / 255.0f - 0.1307f) / 0.3081f; // Normalize to [0, 1] range
            }
        }
    }
#endif
    // Clean up
    stbi_image_free(image_data);
    // Set input tensor with the image data
    infer_request.set_input_tensor(ov::Tensor(ov::element::f32, input_shape, input_data));

    // Perform inference
    infer_request.infer();

    // Get output tensor
    auto output_tensor = infer_request.get_output_tensor();
    auto output_shape = output_tensor.get_shape();
    std::cout << "Output tensor shape: " << output_shape << std::endl;
    auto output_data = output_tensor.data<float>();
    for (size_t i = 0; i < output_tensor.get_size(); ++i) {
        std::cout << output_data[i] << ", ";
    }
    std::cout << std::endl;

    int argmax = std::distance(output_data, std::max_element(output_data, output_data + output_tensor.get_size()));
    std::cout << "Inference result (argmax): " << argmax << std::endl;

    return 0;
}
